from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from typing import Dict, List
import json
import asyncio
import logging
from datetime import datetime

from app.dependencies_json import get_current_user_ws
from app.storage.json_storage import storage

logger = logging.getLogger(__name__)

router = APIRouter()

class ConnectionManager:
    """WebSocket连接管理器"""
    
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.user_connections: Dict[str, str] = {}  # user_id -> connection_id
    
    async def connect(self, websocket: WebSocket, connection_id: str, user_id: str):
        """建立连接"""
        await websocket.accept()
        self.active_connections[connection_id] = websocket
        self.user_connections[user_id] = connection_id
        logger.info(f"WebSocket connected: {connection_id} for user {user_id}")
    
    def disconnect(self, connection_id: str, user_id: str):
        """断开连接"""
        if connection_id in self.active_connections:
            del self.active_connections[connection_id]
        if user_id in self.user_connections:
            del self.user_connections[user_id]
        logger.info(f"WebSocket disconnected: {connection_id} for user {user_id}")
    
    async def send_personal_message(self, message: dict, user_id: str):
        """发送个人消息"""
        connection_id = self.user_connections.get(user_id)
        if connection_id and connection_id in self.active_connections:
            websocket = self.active_connections[connection_id]
            try:
                await websocket.send_text(json.dumps(message))
            except Exception as e:
                logger.error(f"Failed to send message to {user_id}: {e}")
                self.disconnect(connection_id, user_id)
    
    async def broadcast(self, message: dict):
        """广播消息"""
        disconnected = []
        for connection_id, websocket in self.active_connections.items():
            try:
                await websocket.send_text(json.dumps(message))
            except Exception as e:
                logger.error(f"Failed to broadcast to {connection_id}: {e}")
                disconnected.append(connection_id)
        
        # 清理断开的连接
        for connection_id in disconnected:
            for user_id, conn_id in list(self.user_connections.items()):
                if conn_id == connection_id:
                    self.disconnect(connection_id, user_id)
                    break

manager = ConnectionManager()

@router.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    token: str = Query(..., description="访问令牌")
):
    """WebSocket连接端点"""
    connection_id = f"conn_{datetime.utcnow().timestamp()}"
    user = None
    
    try:
        # 验证用户身份
        user = get_current_user_ws(websocket, token)
        user_id = user['id']
        
        # 建立连接
        await manager.connect(websocket, connection_id, user_id)
        
        # 发送欢迎消息
        welcome_message = {
            "type": "welcome",
            "message": f"欢迎 {user['username']}！",
            "timestamp": datetime.utcnow().isoformat(),
            "user_id": user_id
        }
        await websocket.send_text(json.dumps(welcome_message))
        
        # 发送当前系统状态
        await send_system_status(websocket, user_id)
        
        # 监听消息
        while True:
            try:
                data = await websocket.receive_text()
                message = json.loads(data)
                await handle_websocket_message(websocket, message, user)
            except WebSocketDisconnect:
                break
            except json.JSONDecodeError:
                error_message = {
                    "type": "error",
                    "message": "无效的JSON格式",
                    "timestamp": datetime.utcnow().isoformat()
                }
                await websocket.send_text(json.dumps(error_message))
            except Exception as e:
                logger.error(f"WebSocket message handling error: {e}")
                error_message = {
                    "type": "error",
                    "message": "消息处理失败",
                    "timestamp": datetime.utcnow().isoformat()
                }
                await websocket.send_text(json.dumps(error_message))
    
    except Exception as e:
        logger.error(f"WebSocket connection error: {e}")
        try:
            await websocket.close(code=1008)
        except:
            pass
    
    finally:
        if user:
            manager.disconnect(connection_id, user['id'])

async def handle_websocket_message(websocket: WebSocket, message: dict, user: dict):
    """处理WebSocket消息"""
    message_type = message.get("type")
    
    if message_type == "ping":
        # 心跳检测
        pong_message = {
            "type": "pong",
            "timestamp": datetime.utcnow().isoformat()
        }
        await websocket.send_text(json.dumps(pong_message))
    
    elif message_type == "get_tasks":
        # 获取任务列表
        tasks = storage.get_tasks()
        task_message = {
            "type": "tasks_update",
            "data": tasks,
            "timestamp": datetime.utcnow().isoformat()
        }
        await websocket.send_text(json.dumps(task_message))
    
    elif message_type == "get_system_status":
        # 获取系统状态
        await send_system_status(websocket, user['id'])
    
    elif message_type == "subscribe_task":
        # 订阅任务更新
        task_id = message.get("task_id")
        if task_id:
            # 这里可以实现任务订阅逻辑
            response = {
                "type": "subscription_confirmed",
                "task_id": task_id,
                "message": f"已订阅任务 {task_id} 的更新",
                "timestamp": datetime.utcnow().isoformat()
            }
            await websocket.send_text(json.dumps(response))
    
    else:
        # 未知消息类型
        error_message = {
            "type": "error",
            "message": f"未知的消息类型: {message_type}",
            "timestamp": datetime.utcnow().isoformat()
        }
        await websocket.send_text(json.dumps(error_message))

async def send_system_status(websocket: WebSocket, user_id: str):
    """发送系统状态"""
    try:
        # 获取基本统计
        users_count = len(storage.get_users())
        tasks = storage.get_tasks()
        tasks_count = len(tasks)
        
        # 统计任务状态
        status_counts = {}
        for task in tasks:
            status = task.get('status', 'unknown')
            status_counts[status] = status_counts.get(status, 0) + 1
        
        system_status = {
            "type": "system_status",
            "data": {
                "users_count": users_count,
                "tasks_count": tasks_count,
                "task_status_counts": status_counts,
                "storage_type": "JSON Files"
            },
            "timestamp": datetime.utcnow().isoformat()
        }
        
        await websocket.send_text(json.dumps(system_status))
    
    except Exception as e:
        logger.error(f"Failed to send system status: {e}")

# 任务状态更新通知函数（可以被其他模块调用）
async def notify_task_update(task_id: str, task_data: dict):
    """通知任务更新"""
    message = {
        "type": "task_update",
        "task_id": task_id,
        "data": task_data,
        "timestamp": datetime.utcnow().isoformat()
    }
    await manager.broadcast(message)

async def notify_system_event(event_type: str, message: str, data: dict = None):
    """通知系统事件"""
    notification = {
        "type": "system_event",
        "event_type": event_type,
        "message": message,
        "data": data or {},
        "timestamp": datetime.utcnow().isoformat()
    }
    await manager.broadcast(notification)
